import asyncio
from aiohttp import web
import json
import time
import os

# 配置参数（建议通过环境变量或配置文件加载）
VALID_TOKENS=os.environ.get('API_TOKENS', 'sk-123456').split(",")

async def stream_execute_code(code: str, timeout: int = 5):
    """改进的流式执行函数"""
    try:
        # 启动子进程
        proc = await asyncio.create_subprocess_exec(
            "python3", "-u", "-c", code,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )

        # 创建异步队列用于收集输出
        output_queue = asyncio.Queue()

        async def read_stream(stream, is_error=False):
            """异步读取输出流并放入队列"""
            while True:
                line = await stream.readline()
                if not line:
                    break
                await output_queue.put({
                    "content": line.decode().strip(),
                    "is_error": is_error
                })
            await output_queue.put(None)  # 结束标记

        # 启动两个读取任务
        readers = [
            asyncio.create_task(read_stream(proc.stdout)),
            asyncio.create_task(read_stream(proc.stderr, True))
        ]

        start_time = asyncio.get_event_loop().time()
        active_readers = 2

        while True:
            # 处理超时
            if (asyncio.get_event_loop().time() - start_time) > timeout:
                proc.kill()
                yield {
                    "content": f"\nExecution timed out after {timeout} seconds",
                    "is_error": True
                }
                break

            # 从队列获取输出
            try:
                item = await asyncio.wait_for(output_queue.get(), timeout=0.1)
                if item is None:
                    active_readers -= 1
                    if active_readers == 0:
                        break
                    continue
                yield item
            except asyncio.TimeoutError:
                # 检查进程是否已结束
                if proc.returncode is not None:
                    break
                continue

        # 清理任务
        for task in readers:
            task.cancel()

    except Exception as e:
        yield {"content": str(e), "is_error": True}


async def token_auth_middleware(app, handler):
    """Token验证中间件"""
    async def middleware_handler(request):
        # 跳过预检请求
        if request.method == "OPTIONS":
            return await handler(request)
        
        # 获取Authorization头
        auth_header = request.headers.get("Authorization", "")
        
        # 验证格式
        if not auth_header.startswith("Bearer "):
            return web.Response(
                status=401,
                text=json.dumps({"error": "Invalid authorization format"}),
                content_type="application/json"
            )
        
        # 提取token
        token = auth_header[7:].strip()
        
        # 验证token有效性
        if token not in VALID_TOKENS:
            return web.Response(
                status=403,
                text=json.dumps({"error": "Invalid API token"}),
                content_type="application/json"
            )
        
        return await handler(request)
    return middleware_handler

async def handle_stream_request(request: web.Request) -> web.StreamResponse:
    response = web.StreamResponse(
        headers={
            'Content-Type': 'text/event-stream',
            'Access-Control-Allow-Origin': '*',  # CORS设置
            'Access-Control-Allow-Headers': 'Authorization'  # 允许的请求头
        }
    )
    await response.prepare(request)

    try:
        # 解析请求数据
        data = await request.json()
        code = next(
            (msg["content"] for msg in reversed(data.get("messages", [])) 
             if msg["role"] == "user"),
            ""
        )

        if not code:
            await response.write(b'data: {"error": "No code provided"}\n\n')
            return response

        # 执行代码流
        has_output = False
        async for chunk in stream_execute_code(code):

            # 构造响应数据
            content = chunk["content"]
            if chunk["is_error"]:
                content = f"[ERROR] {content}" if content else ""

            event_data = {
                "id": "chatcmpl-123",
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": "code-executor",
                "choices": [{
                    "delta": {"content": content + "\n"},
                    "index": 0,
                    "finish_reason": None
                }]
            }

            await response.write(
                f"data: {json.dumps(event_data)}\n\n".encode()
            )
            has_output = True

        # 发送结束标记
        await response.write(b'data: [DONE]\n\n')

    except Exception as e:
        error_data = {
            "error": str(e),
            "code": 500
        }
        await response.write(
            f"data: {json.dumps(error_data)}\n\n".encode()
        )
    finally:
        await response.write_eof()
    
    return response

# 创建应用时添加中间件
app = web.Application(middlewares=[token_auth_middleware])
app.router.add_post("/v1/chat/completions", handle_stream_request)
app.router.add_options("/v1/chat/completions", lambda r: web.Response())  # CORS预检支持

if __name__ == "__main__":
    web.run_app(app, port=8080, host='0.0.0.0')
